import sys
import os
import cv2
import numpy as np
import json
from PyQt6.QtWidgets import (QApplication, QWidget, QLabel, QLineEdit, QPushButton,
                             QVBoxLayout, QHBoxLayout, QSlider, QCheckBox,
                             QMainWindow, QFrame, QGridLayout, QMessageBox, QGroupBox)
from PyQt6.QtGui import QImage, QPixmap, QIcon, QGuiApplication, QIntValidator, QPainter, QColor
from PyQt6.QtCore import Qt, QThread, pyqtSignal, QSize

def resource_path(relative_path):
    """ 获取资源的绝对路径，无论是从脚本运行还是从打包后的exe运行 """
    try:
        # PyInstaller 创建一个临时文件夹，并将路径存储在 _MEIPASS 中
        base_path = sys._MEIPASS
    except Exception:
        base_path = os.path.abspath(".")

    return os.path.join(base_path, relative_path)

# --- Unicorn主题UI样式表 (QSS) ---
# 注意：这里不能使用相对路径，所以背景图将在代码中设置
UNICORN_STYLE_SHEET = """
QWidget {
    color: #333;
    font-family: 'Segoe UI', 'Microsoft YaHei', sans-serif;
}

QMainWindow {
    background-color: #e9e8f5; /* 从图片提取的柔和背景色 */
}

QGroupBox {
    font-size: 14px;
    font-weight: bold;
    color: #8e7cc3; /* 从图片提取的淡紫色 */
    border: 1px solid rgba(142, 124, 195, 0.3);
    border-radius: 8px;
    margin-top: 10px;
    background-color: rgba(255, 255, 255, 0.7); /* 半透明白色背景，突出内容 */
}

QGroupBox::title {
    subcontrol-origin: margin;
    subcontrol-position: top center;
    padding: 0 10px;
    background-color: rgba(233, 232, 245, 0.8);
    border-radius: 4px;
}

QLabel {
    color: #555;
    font-size: 12px;
    background-color: transparent; /* 确保标签背景透明 */
}

QLineEdit {
    background-color: rgba(255, 255, 255, 0.8);
    border: 1px solid #dcdcf0;
    border-radius: 4px;
    padding: 5px;
    color: #333;
}

QLineEdit:focus {
    border: 1px solid #8e7cc3;
}

QPushButton {
    background-color: #8e7cc3;
    color: white;
    border: none;
    padding: 8px 16px;
    border-radius: 4px;
    font-size: 13px;
    font-weight: bold;
}

QPushButton:hover {
    background-color: #a08eda;
}

QPushButton:pressed {
    background-color: #7b68b0;
}

QPushButton:disabled {
    background-color: #cccccc;
    color: #888888;
}

QCheckBox {
    spacing: 10px;
    background-color: transparent;
}

QCheckBox::indicator {
    width: 18px;
    height: 18px;
    border-radius: 4px;
}

QCheckBox::indicator:unchecked {
    background-color: #e0e0f0;
    border: 1px solid #c0c0d8;
}

QCheckBox::indicator:checked {
    background-color: #8e7cc3;
}

QSlider::groove:horizontal {
    border: 1px solid #dcdcf0;
    height: 4px;
    background: #dcdcf0;
    margin: 2px 0;
    border-radius: 2px;
}

QSlider::handle:horizontal {
    background: #8e7cc3;
    border: 1px solid #8e7cc3;
    width: 16px;
    height: 16px;
    margin: -6px 0;
    border-radius: 8px;
}

QSlider::sub-page:horizontal {
    background: #a08eda;
    border: 1px solid #dcdcf0;
    height: 4px;
    border-radius: 2px;
}
"""

# --- 视频处理线程 ---
class VideoThread(QThread):
    change_pixmap_signal = pyqtSignal(np.ndarray)
    status_signal = pyqtSignal(float, float, bool)
    connection_failed_signal = pyqtSignal()
    stream_finished_signal = pyqtSignal()

    def __init__(self, source_url, filters):
        super().__init__()
        self._run_flag = True
        self.source_url = source_url
        self.filters = filters

    def run(self):
        cap = cv2.VideoCapture(self.source_url, cv2.CAP_FFMPEG)
        cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)

        if not cap.isOpened():
            print(f"错误: 无法打开视频流 {self.source_url}")
            self.connection_failed_signal.emit()
            return

        frame_count = 0
        start_time = cv2.getTickCount()

        while self._run_flag:
            latency_start = cv2.getTickCount()
            ret, cv_img = cap.read()
            latency_end = cv2.getTickCount()
            latency_ms = (latency_end - latency_start) / cv2.getTickFrequency() * 1000

            if ret:
                if self.filters.get('sharpen_enabled'):
                    strength = self.filters.get('sharpen_strength', 0.0)
                    center_val = 9 + strength * 4
                    kernel = np.array([[-1, -1, -1], [-1, center_val, -1], [-1, -1, -1]])
                    cv_img = cv2.filter2D(cv_img, -1, kernel)

                if self.filters.get('denoise_enabled'):
                    strength = self.filters.get('denoise_strength', 0)
                    if strength > 0:
                        cv_img = cv2.fastNlMeansDenoisingColored(cv_img, None, float(strength), 10, 7, 21)

                frame_count += 1
                elapsed_time = (cv2.getTickCount() - start_time) / cv2.getTickFrequency()
                fps = frame_count / elapsed_time if elapsed_time > 0 else 0

                self.change_pixmap_signal.emit(cv_img)
                self.status_signal.emit(fps, latency_ms, True)
            else:
                print("视频流已结束或连接超时。")
                self.stream_finished_signal.emit()
                break

        cap.release()
        self.status_signal.emit(0.0, 0.0, False)

    def stop(self):
        self._run_flag = False
        self.wait()
    
    def update_filters(self, new_filters):
        self.filters = new_filters

# --- 独立的播放器窗口 ---
class PlayerWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("视频播放")
        self.video_label = QLabel(self)
        self.video_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.video_label.setStyleSheet("background-color: black;")
        
        layout = QVBoxLayout()
        layout.setContentsMargins(0,0,0,0)
        layout.addWidget(self.video_label)
        self.setLayout(layout)

        self.setWindowFlags(Qt.WindowType.FramelessWindowHint)

    def update_frame(self, frame):
        qt_img = self.convert_cv_qt(frame)
        self.video_label.setPixmap(qt_img)

    def convert_cv_qt(self, cv_img):
        rgb_image = cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB)
        h, w, ch = rgb_image.shape
        bytes_per_line = ch * w
        convert_to_Qt_format = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format.Format_RGB888)
        return QPixmap.fromImage(convert_to_Qt_format)

    def resize_and_center(self, w, h):
        self.setFixedSize(w, h)
        screen_geometry = QGuiApplication.primaryScreen().geometry()
        x = (screen_geometry.width() - self.width()) // 2
        y = (screen_geometry.height() - self.height()) // 2
        self.move(x, y)

    def set_always_on_top(self, on_top):
        is_visible = self.isVisible()
        if on_top:
            self.setWindowFlags(self.windowFlags() | Qt.WindowType.WindowStaysOnTopHint)
        else:
            self.setWindowFlags(self.windowFlags() & ~Qt.WindowType.WindowStaysOnTopHint)
        if is_visible:
            self.show()

# --- 带背景图的自定义控件 ---
class BackgroundWidget(QWidget):
    def __init__(self, image_path, parent=None):
        super().__init__(parent)
        self.background_image = QPixmap(image_path)
    
    def paintEvent(self, event):
        painter = QPainter(self)
        painter.drawPixmap(self.rect(), self.background_image.scaled(
            self.size(), 
            Qt.AspectRatioMode.KeepAspectRatioByExpanding, 
            Qt.TransformationMode.SmoothTransformation
        ))
        painter.fillRect(self.rect(), QColor(255, 255, 255, 120))


# --- 主控制窗口 ---
class ControlWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Unicorn UDP 视频接收器")
        self.setFixedSize(380, 480)
        
        # 修复：使用 resource_path 函数加载图标
        self.setWindowIcon(QIcon(resource_path("unicorn.jpg")))

        # 修复：使用 resource_path 函数加载背景图
        self.central_widget = BackgroundWidget(resource_path("unicorn.jpg"))
        self.setCentralWidget(self.central_widget)
        
        self.main_layout = QVBoxLayout(self.central_widget)
        self.main_layout.setContentsMargins(15, 15, 15, 15)
        self.main_layout.setSpacing(10)

        self.video_thread = None
        self.player_window = PlayerWindow()
        
        # 修复：定义配置文件路径，使其与exe在同一目录
        self.config_file_path = self._get_config_path()

        self.init_ui()
        self.load_config()

    def _get_config_path(self):
        """获取配置文件的正确路径，使其与exe或py文件在同一目录"""
        if getattr(sys, 'frozen', False):
            # 如果是打包后的exe
            application_path = os.path.dirname(sys.executable)
        else:
            # 如果是直接运行的py脚本
            application_path = os.path.dirname(os.path.abspath(__file__))
        return os.path.join(application_path, "config.json")

    def init_ui(self):
        # UI初始化代码... (与之前版本相同)
        settings_group = QGroupBox("网络与分辨率设置")
        grid_layout = QGridLayout(settings_group)
        grid_layout.addWidget(QLabel("IP地址:"), 0, 0)
        self.ip_input = QLineEdit("0.0.0.0")
        grid_layout.addWidget(self.ip_input, 0, 1, 1, 3)
        grid_layout.addWidget(QLabel("端口:"), 1, 0)
        self.port_input = QLineEdit("60006")
        self.port_input.setValidator(QIntValidator(1024, 65535, self))
        grid_layout.addWidget(self.port_input, 1, 1)
        grid_layout.addWidget(QLabel("宽度:"), 2, 0)
        self.width_input = QLineEdit("640")
        self.width_input.setValidator(QIntValidator(1, 7680, self))
        grid_layout.addWidget(self.width_input, 2, 1)
        grid_layout.addWidget(QLabel("高度:"), 2, 2)
        self.height_input = QLineEdit("480")
        self.height_input.setValidator(QIntValidator(1, 4320, self))
        grid_layout.addWidget(self.height_input, 2, 3)
        self.main_layout.addWidget(settings_group)
        processing_group = QGroupBox("画质增强")
        proc_layout = QVBoxLayout(processing_group)
        sharpen_layout = QHBoxLayout()
        self.sharpen_check = QCheckBox("启用锐化")
        self.sharpen_slider = QSlider(Qt.Orientation.Horizontal)
        self.sharpen_slider.setRange(0, 100)
        self.sharpen_label = QLabel("强度: 0.00")
        sharpen_layout.addWidget(self.sharpen_check)
        sharpen_layout.addWidget(self.sharpen_slider)
        sharpen_layout.addWidget(self.sharpen_label)
        denoise_layout = QHBoxLayout()
        self.denoise_check = QCheckBox("启用降噪")
        self.denoise_check.setToolTip("警告：此功能会极大地影响帧率和延迟！")
        self.denoise_slider = QSlider(Qt.Orientation.Horizontal)
        self.denoise_slider.setRange(0, 20)
        self.denoise_label = QLabel("强度: 0")
        denoise_layout.addWidget(self.denoise_check)
        denoise_layout.addWidget(self.denoise_slider)
        denoise_layout.addWidget(self.denoise_label)
        proc_layout.addLayout(sharpen_layout)
        proc_layout.addLayout(denoise_layout)
        self.main_layout.addWidget(processing_group)
        control_group = QGroupBox("控制")
        control_layout = QHBoxLayout(control_group)
        self.start_button = QPushButton("开始接收")
        self.stop_button = QPushButton("停止接收")
        self.save_button = QPushButton("保存配置")
        self.stop_button.setEnabled(False)
        control_layout.addWidget(self.start_button)
        control_layout.addWidget(self.stop_button)
        control_layout.addWidget(self.save_button)
        self.main_layout.addWidget(control_group)
        status_group = QGroupBox("状态与选项")
        status_layout = QVBoxLayout(status_group)
        self.always_on_top_check = QCheckBox("播放窗口置顶")
        status_layout.addWidget(self.always_on_top_check)
        status_line_layout = QHBoxLayout()
        self.fps_label = QLabel("帧率: -- fps")
        self.latency_label = QLabel("延迟: -- ms")
        status_line_layout.addWidget(self.fps_label)
        status_line_layout.addWidget(self.latency_label)
        status_line_layout.addStretch()
        self.connection_label = QLabel("状态: 未连接")
        status_line_layout.addWidget(self.connection_label)
        status_layout.addLayout(status_line_layout)
        self.main_layout.addWidget(status_group)
        self.main_layout.addStretch()
        self.start_button.clicked.connect(self.start_receiving)
        self.stop_button.clicked.connect(self.stop_receiving)
        self.save_button.clicked.connect(self.save_config)
        self.sharpen_check.stateChanged.connect(self._update_filters_in_thread)
        self.denoise_check.stateChanged.connect(self._update_filters_in_thread)
        self.sharpen_slider.valueChanged.connect(self._update_filters_in_thread)
        self.denoise_slider.valueChanged.connect(self._update_filters_in_thread)
        self.always_on_top_check.stateChanged.connect(self._update_always_on_top)

    def _update_filters_in_thread(self):
        self.sharpen_label.setText(f"强度: {self.sharpen_slider.value()/100.0:.2f}")
        self.denoise_label.setText(f"强度: {self.denoise_slider.value()}")
        if self.video_thread and self.video_thread.isRunning():
            filters = {
                'sharpen_enabled': self.sharpen_check.isChecked(),
                'sharpen_strength': self.sharpen_slider.value() / 100.0,
                'denoise_enabled': self.denoise_check.isChecked(),
                'denoise_strength': self.denoise_slider.value()
            }
            self.video_thread.update_filters(filters)

    def _update_always_on_top(self):
        if self.player_window:
            self.player_window.set_always_on_top(self.always_on_top_check.isChecked())

    def start_receiving(self):
        try:
            ip = self.ip_input.text()
            port = self.port_input.text()
            width = int(self.width_input.text())
            height = int(self.height_input.text())
            if not all([ip, port, self.width_input.text(), self.height_input.text()]):
                raise ValueError("输入字段不能为空")
            source_url = f"udp://{ip}:{port}?overrun_nonfatal=1&fifo_size=100000&timeout=5000000"
            initial_filters = {
                'sharpen_enabled': self.sharpen_check.isChecked(),
                'sharpen_strength': self.sharpen_slider.value() / 100.0,
                'denoise_enabled': self.denoise_check.isChecked(),
                'denoise_strength': self.denoise_slider.value()
            }
            self.video_thread = VideoThread(source_url, initial_filters)
            self.video_thread.change_pixmap_signal.connect(self.player_window.update_frame)
            self.video_thread.status_signal.connect(self.update_status)
            self.video_thread.connection_failed_signal.connect(self.handle_connection_failure)
            self.video_thread.stream_finished_signal.connect(self.handle_stream_finished)
            self.video_thread.start()
            self.player_window.set_always_on_top(self.always_on_top_check.isChecked())
            self.player_window.resize_and_center(width, height)
            self.player_window.show()
            self.start_button.setEnabled(False)
            self.stop_button.setEnabled(True)
        except ValueError:
            self.show_input_error_message()

    def stop_receiving(self):
        if self.video_thread:
            self.video_thread.stop()
            self.video_thread = None
        self.reset_ui_on_stop()

    def reset_ui_on_stop(self):
        self.player_window.hide()
        self.start_button.setEnabled(True)
        self.stop_button.setEnabled(False)
        self.update_status(0.0, 0.0, False)

    def handle_connection_failure(self):
        msg_box = QMessageBox(self)
        msg_box.setIcon(QMessageBox.Icon.Critical)
        msg_box.setText("无法连接到视频流！")
        msg_box.setInformativeText("请检查:\n1. IP地址和端口号是否正确。\n2. OBS推流是否已开启。\n3. 防火墙设置。")
        msg_box.setWindowTitle("连接失败")
        msg_box.exec()
        self.reset_ui_on_stop()
        
    def handle_stream_finished(self):
        print("接收到视频流结束信号，正在关闭播放窗口并重置UI。")
        self.stop_receiving()

    def show_input_error_message(self):
        msg_box = QMessageBox(self)
        msg_box.setIcon(QMessageBox.Icon.Warning)
        msg_box.setText("输入参数无效！")
        msg_box.setInformativeText("请检查所有输入字段是否填写正确。")
        msg_box.setWindowTitle("输入错误")
        msg_box.exec()

    def update_status(self, fps, latency, is_connected):
        self.fps_label.setText(f"帧率: {fps:.2f} fps")
        self.latency_label.setText(f"延迟: {latency:.1f} ms")
        if is_connected:
            self.connection_label.setText(f"状态: <font color='#8e7cc3'>已连接</font>")
        else:
            self.connection_label.setText("状态: <font color='#e74c3c'>未连接</font>")

    def save_config(self):
        config = {
            "ip": self.ip_input.text(),
            "port": self.port_input.text(),
            "width": self.width_input.text(),
            "height": self.height_input.text(),
            "sharpen_enabled": self.sharpen_check.isChecked(),
            "sharpen_strength": self.sharpen_slider.value(),
            "denoise_enabled": self.denoise_check.isChecked(),
            "denoise_strength": self.denoise_slider.value(),
            "always_on_top": self.always_on_top_check.isChecked()
        }
        try:
            with open(self.config_file_path, 'w') as f:
                json.dump(config, f, indent=4)
            print("配置已保存。")
        except Exception as e:
            print(f"保存配置失败: {e}")

    def load_config(self):
        try:
            with open(self.config_file_path, 'r') as f:
                config = json.load(f)
                self.ip_input.setText(config.get("ip", "0.0.0.0"))
                self.port_input.setText(config.get("port", "60006"))
                self.width_input.setText(config.get("width", "640"))
                self.height_input.setText(config.get("height", "480"))
                self.sharpen_check.setChecked(config.get("sharpen_enabled", False))
                self.sharpen_slider.setValue(config.get("sharpen_strength", 0))
                self.denoise_check.setChecked(config.get("denoise_enabled", False))
                self.denoise_slider.setValue(config.get("denoise_strength", 0))
                self.always_on_top_check.setChecked(config.get("always_on_top", False))
                self._update_filters_in_thread()
                print("配置已加载。")
        except FileNotFoundError:
            print("未找到配置文件，使用默认设置。")
        except Exception as e:
            print(f"加载配置失败: {e}")

    def closeEvent(self, event):
        self.stop_receiving()
        self.player_window.close()
        event.accept()

def main():
    app = QApplication(sys.argv)
    # 修复：不再使用QSS设置背景图，因为 paintEvent 会处理
    # app.setStyleSheet(UNICORN_STYLE_SHEET) 
    window = ControlWindow()
    window.setStyleSheet(UNICORN_STYLE_SHEET) # 将样式表应用到主窗口
    window.show()
    sys.exit(app.exec())

if __name__ == "__main__":
    main()
